from pymongo import MongoClient
from pymongo.errors import PyMongoError
from typing import Union, Dict, Any, List, Optional
import pandas as pd

# nebula_data_engine
class MongoDataEngine:
    def __init__(
        self,
        conn_str: str,
        table_name: Optional[str] = None,
        return_format: str = "json",
        auto_create_table: bool = True,
    ):
        self.conn_str = conn_str
        self.return_format = return_format
        self.auto_create_table = auto_create_table

        # 解析连接字符串，提取数据库名称
        self.client = MongoClient(conn_str)
        self.db_name = self._extract_db_name(conn_str)
        self.db = self.client[self.db_name]

        # 默认集合名
        self.default_table_name = table_name or "default"  # 如果未指定表名，则使用 "default"

        # 初始化集合
        self.tables = {}

    def _extract_db_name(self, conn_str: str) -> str:
        """
        从连接字符串中提取数据库名称。
        如果未指定数据库名称，则使用默认数据库 "test"。
        """
        if "/" not in conn_str:
            return "test"  # 默认数据库
        db_part = conn_str.split("/")[-1]  # 获取数据库部分
        if "?" in db_part:
            db_part = db_part.split("?")[0]  # 去掉查询参数
        return db_part or "test"  # 如果数据库名称为空，则使用 "test"

    def _get_collection(self, table_name: Optional[str] = None):
        """
        获取集合对象。
        如果未指定集合名，则使用默认集合名。
        """
        if table_name is None:
            table_name = self.default_table_name

        if table_name not in self.tables:
            if self.auto_create_table:
                self.tables[table_name] = self.db[table_name]
            else:
                raise ValueError(f"集合 '{table_name}' 不存在，且自动建表已关闭。")

        return self.tables[table_name]

    def save(self, entity: Dict[str, Any], table_name: Optional[str] = None) -> str:
        """
        插入数据。
        """
        collection = self._get_collection(table_name)
        result = collection.insert_one(entity)
        return str(result.inserted_id)

    def delete(self, query: Dict[str, Any], table_name: Optional[str] = None) -> int:
        """
        删除数据。
        """
        collection = self._get_collection(table_name)
        result = collection.delete_many(query)
        return result.deleted_count

    def deleteById(self, id: str, table_name: Optional[str] = None) -> int:
        """
        根据 ID 删除数据。
        """
        collection = self._get_collection(table_name)
        result = collection.delete_one({"_id": id})
        return result.deleted_count

    def update(self, query: Dict[str, Any], update_data: Dict[str, Any], table_name: Optional[str] = None) -> int:
        """
        更新数据。
        """
        collection = self._get_collection(table_name)
        result = collection.update_many(query, {"$set": update_data})
        return result.modified_count

    def updateById(self, id: str, update_data: Dict[str, Any], table_name: Optional[str] = None) -> int:
        """
        根据 ID 更新数据。
        """
        collection = self._get_collection(table_name)
        result = collection.update_one({"_id": id}, {"$set": update_data})
        return result.modified_count

    def findOneAndUpdate(
        self,
        query: Dict[str, Any],
        update_data: Dict[str, Any],
        return_document: bool = False,
        table_name: Optional[str] = None,
    ) -> Union[Dict, None]:
        """
        查询并更新一条数据。
        :param query: 查询条件。
        :param update_data: 更新数据。
        :param return_document: 是否返回更新后的文档（True 返回更新后，False 返回更新前）。
        :param table_name: 集合名称。
        :return: 更新前或更新后的文档。
        """
        collection = self._get_collection(table_name)
        result = collection.find_one_and_update(
            query,
            {"$set": update_data},
            return_document=True if return_document else False,
        )
        return result

    def findByIdAndUpdate(
        self,
        id: str,
        update_data: Dict[str, Any],
        return_document: bool = False,
        table_name: Optional[str] = None,
    ) -> Union[Dict, None]:
        """
        根据 ID 查询并更新一条数据。
        :param id: 文档 ID。
        :param update_data: 更新数据。
        :param return_document: 是否返回更新后的文档（True 返回更新后，False 返回更新前）。
        :param table_name: 集合名称。
        :return: 更新前或更新后的文档。
        """
        collection = self._get_collection(table_name)
        result = collection.find_one_and_update(
            {"_id": id},
            {"$set": update_data},
            return_document=True if return_document else False,
        )
        return result

    def find(self, query: Dict[str, Any], table_name: Optional[str] = None) -> Union[List[Dict], pd.DataFrame]:
        """
        查询数据。
        """
        collection = self._get_collection(table_name)
        documents = list(collection.find(query))  # 包含 _id 字段

        if self.return_format == "json":
            return documents
        else:
            return pd.DataFrame(documents)

    def findBy(self, query: Dict[str, Any], table_name: Optional[str] = None) -> Union[List[Dict], pd.DataFrame]:
        """
        根据条件查询数据。
        """
        return self.find(query, table_name)

    def findAll(self, table_name: Optional[str] = None) -> Union[List[Dict], pd.DataFrame]:
        """
        查询集合中的所有数据。
        """
        collection = self._get_collection(table_name)
        documents = list(collection.find({}))  # 包含 _id 字段

        if self.return_format == "json":
            return documents
        else:
            return pd.DataFrame(documents)

    def findById(self, id: str, table_name: Optional[str] = None) -> Union[Dict, pd.DataFrame]:
        """
        根据 ID 查询数据。
        """
        collection = self._get_collection(table_name)
        document = collection.find_one({"_id": id})  # 包含 _id 字段

        if self.return_format == "json":
            return document
        else:
            return pd.DataFrame([document] if document else [])

    def findPage(self, query: Dict[str, Any], page: int, page_size: int, table_name: Optional[str] = None) -> Union[List[Dict], pd.DataFrame]:
        """
        分页查询数据。
        """
        collection = self._get_collection(table_name)
        skip = (page - 1) * page_size
        documents = list(collection.find(query).skip(skip).limit(page_size))  # 包含 _id 字段

        if self.return_format == "json":
            return documents
        else:
            return pd.DataFrame(documents)

    def transaction(self, operations: List[Dict[str, Any]]) -> bool:
        """
        执行事务操作。
        operations 格式：
        [
            {"action": "save", "table_name": "users", "data": {"name": "Alice", "age": 25}},
            {"action": "update", "table_name": "users", "query": {"name": "Alice"}, "update_data": {"age": 26}},
            {"action": "delete", "table_name": "users", "query": {"name": "Bob"}}
        ]
        """
        try:
            with self.client.start_session() as session:
                with session.start_transaction():
                    for op in operations:
                        action = op.get("action")
                        table_name = op.get("table_name")
                        if action == "save":
                            self.save(op["data"], table_name)
                        elif action == "update":
                            self.update(op["query"], op["update_data"], table_name)
                        elif action == "delete":
                            self.delete(op["query"], table_name)
                        else:
                            raise ValueError(f"不支持的操作: {action}")
                    return True
        except PyMongoError as e:
            print(f"事务执行失败: {e}")
            return False